"""
简易推理脚本：从纯噪声生成 Hi-C 窗口（基于已训练 Squidiff 模型）

用法示例：
python infer_generate_hic.py \
  --checkpoint_path model/uncond_checkpoints \
  --output_dir result/infer \
  --timepoints 0.25,0.5,0.75 \
  --window_size 256 \
  --num_samples 1 \
  --ddim_steps 50 \
  --use_fp16

核心流程：
1. 解析 CLI 参数 → 获取时间点、窗口大小、样本数等。
2. dist_util.setup_dist() 初始化显卡/多进程。
3. load_model_from_checkpoint() 载入 UNet + diffusion；prefer EMA 参数。
4. 对每个时间点 t：
   a. 构造生物学时间张量 bio_time=[t]*B。
   b. diffusion.p_sample_loop(...) 传入 shape=(B,1,H,W)、model_kwargs={bio_time}，
      使用 DDIM 或全扩散步骤，从纯噪声生成窗口张量 (B,1,H,W)。
   c. 保存 arr.float32 -> .npy；对 arr.mean() 生成 PNG（log1p+分位裁剪）。

说明：
- 该脚本不依赖数据集，完全从噪声产生；适用于 unconditional / time-conditional 模型。
- 若模型需要额外语义/条件，请在 model_kwargs 中补充。
"""

import os
import argparse
import numpy as np

from Squidiff import dist_util, logger
from Squidiff.hic_script_util import get_all_defaults, add_dict_to_argparser
from Squidiff.hic_inference import load_model_from_checkpoint


def parse_args():
    defaults = get_all_defaults()
    parser = argparse.ArgumentParser(
        description="Generate Hi-C windows from pure noise (Squidiff model)",
        conflict_handler='resolve'
    )

    # 必要参数
    parser.add_argument("--checkpoint_path", type=str, required=True, help="Checkpoint 文件或目录")
    parser.add_argument("--output_dir", type=str, required=True, help="输出目录")

    # 推理相关
    parser.add_argument("--timepoints", type=str, default="0.5", help="逗号分隔的生物学时间点，如 0.25,0.5")
    parser.add_argument("--window_size", type=int, default=256, help="生成窗口大小 H=W")
    parser.add_argument("--num_samples", type=int, default=1, help="每个时间点生成多少样本")
    parser.add_argument("--ddim_steps", type=int, default=50, help="DDIM 步数；0 表示使用全扩散 p_sample_loop")
    parser.add_argument("--use_fp16", action="store_true", help="半精度推理")

    # 可视化
    parser.add_argument("--use_log", type=str, default="True", help="png 显示 log1p")
    parser.add_argument("--vmin_pct", type=float, default=2.0)
    parser.add_argument("--vmax_pct", type=float, default=98.0)

    # 加载默认模型参数（避免缺项）
    add_dict_to_argparser(parser, defaults)

    args = parser.parse_args()
    args = vars(args)
    args["use_log"] = str(args["use_log"]).lower() in ("1", "true", "t", "y", "yes")
    os.makedirs(args["output_dir"], exist_ok=True)
    return args


def visualize(arr: np.ndarray, save_path: str, use_log: bool, vmin_pct: float, vmax_pct: float):
    try:
        import matplotlib.pyplot as plt
        mat = arr.copy()
        if use_log:
            mat = np.log1p(mat)
        vmin = vmax = None
        if vmin_pct >= 0 and vmax_pct >= 0:
            vmin = np.percentile(mat, vmin_pct)
            vmax = np.percentile(mat, vmax_pct)
            if vmin == vmax:
                vmin = vmax = None
        fig = plt.figure(figsize=(4,4), dpi=150)
        ax = fig.add_subplot(111)
        im = ax.imshow(mat, cmap='Reds', origin='lower', vmin=vmin, vmax=vmax)
        ax.axis('off')
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.tight_layout()
        fig.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close(fig)
    except Exception as e:
        print(f"[WARN] matplotlib failed: {e}")


def main():
    args = parse_args()

    # Step1 设备与日志
    dist_util.setup_dist()
    logger.configure(dir=args["output_dir"])

    # Step2 载入模型+diffusion（prefer EMA）
    model, diffusion, _ = load_model_from_checkpoint(args["checkpoint_path"], args, use_ema=True)

    ws = int(args["window_size"])
    shape = (int(args["num_samples"]), 1, ws, ws)
    device = dist_util.dev()

    timepoints = [float(t) for t in args["timepoints"].split(',') if t.strip()]

    for t in timepoints:
        logger.log(f"Generating {args['num_samples']} window(s) at time {t} ...")
        bio_times = [t] * int(args["num_samples"])

        # 生成
        with np.errstate(over='ignore'):  # numpy log1p 可触发警告
            import torch as th
            if args["ddim_steps"] > 0:
                samples = diffusion.p_sample_loop(
                    model, shape, noise=None, clip_denoised=False,
                    model_kwargs={'bio_time': th.tensor(bio_times, device=device)},
                    device=device, progress=False
                )
            else:
                samples = diffusion.p_sample_loop(
                    model, shape, noise=None, clip_denoised=False,
                    model_kwargs={'bio_time': th.tensor(bio_times, device=device)},
                    device=device, progress=False
                )

        arr = samples.detach().cpu().numpy()[:, 0]  # (B,H,W)

        # 保存 .npy
        npy_path = os.path.join(args["output_dir"], f"hic_t{t:.3f}_fp32.npy")
        np.save(npy_path, arr.astype(np.float32))

        # 保存 PNG（若多张则存均值）
        mat_show = arr.mean(axis=0) if arr.shape[0] > 1 else arr[0]
        png_path = os.path.join(args["output_dir"], f"hic_t{t:.3f}.png")
        visualize(mat_show, png_path, args["use_log"], args["vmin_pct"], args["vmax_pct"])

        logger.log(f"Saved npy -> {npy_path}; png -> {png_path}")


if __name__ == "__main__":
    main()
